import torch
import torch.nn as nn
from model.common import Focus, Conv, CSPconv, SPP, Classify


class CSPDarkNet(nn.Module):
    def __init__(self):
        super().__init__()
        self.focus = Focus()
        self.conv1 = Conv(c1=32, c2=64, k=3, s=2)
        self.csp1 = CSPconv(c1=64, c2=64)
        self.conv2 = Conv(c1=64, c2=128, k=3, s=2)
        self.csp2 = CSPconv(c1=128, c2=128, n=3)
        self.conv3 = Conv(c1=128, c2=256, k=3, s=2)
        self.csp3 = CSPconv(c1=256, c2=256, n=3)
        self.conv4 = Conv(c1=256, c2=512, k=3, s=2)
        self.spp = SPP(512, 512)

    def forward(self, x):
        y = []
        x = self.focus(x)
        x = self.conv1(x)
        x = self.csp1(x)
        x = self.conv2(x)
        x = self.csp2(x)
        y.append(x)
        x = self.conv3(x)
        x = self.csp3(x)
        y.append(x)
        x = self.conv4(x)
        x = self.spp(x)
        y.append(x)
        return y


class Neck(nn.Module):
    def __init__(self):
        super().__init__()
        self.upsample = nn.Upsample(scale_factor=2)
        self.csp1 = CSPconv(c1=512, c2=512, shortcut=False)
        self.conv1 = Conv(c1=512, c2=256)
        self.csp2 = CSPconv(c1=512, c2=256, shortcut=False)
        self.conv2 = Conv(c1=256, c2=128)
        self.csp3 = CSPconv(c1=256, c2=128, shortcut=False)
        self.conv3 = Conv(c1=128, c2=128, k=3, s=2)
        self.csp4 = CSPconv(c1=256, c2=256, shortcut=False)
        self.conv4 = Conv(c1=256, c2=256, k=3, s=2)
        self.csp5 = CSPconv(c1=512, c2=512, shortcut=False)

    def forward(self, x):
        x1, x2, x3 = x
        y3 = self.csp1(x3)
        y3 = self.conv1(y3)
        y2 = self.upsample(y3)
        y2 = torch.cat([y2, x2], dim=1)
        y2 = self.csp2(y2)
        y2 = self.conv2(y2)
        y1 = self.upsample(y2)
        y1 = torch.cat([y1, x1], dim=1)
        z1 = self.csp3(y1)
        z2 = self.conv3(z1)
        z2 = torch.cat([z2, y2], dim=1)
        z2 = self.csp4(z2)
        z3 = self.conv4(z2)
        z3 = torch.cat([z3, y3], dim=1)
        z3 = self.csp5(z3)
        return [z1, z2, z3]


class Head(nn.Module):
    def __init__(self):
        super().__init__()
        self.conv1 = nn.Conv2d(in_channels=128, out_channels=255, kernel_size=1)
        self.conv2 = nn.Conv2d(in_channels=256, out_channels=255, kernel_size=1)
        self.conv3 = nn.Conv2d(in_channels=512, out_channels=255, kernel_size=1)

    def forward(self, x):
        x = [self.conv1(x[0]), self.conv2(x[1]), self.conv3(x[2])]
        return x


class YoloV5(nn.Module):
    def __init__(self, classify=False):
        super().__init__()
        self.backbone = CSPDarkNet()
        self.neck = Neck()
        self.head = Head()
        self.classify = Classify(128 + 256 + 512, 10)
        self.classify_forward = classify

    def forward(self, x):
        x = self.backbone(x)
        if self.classify_forward:
            x = self.classify(x)
        else:
            x = self.neck(x)
            x = self.head(x)
        return x
